
from isaaclab.utils import configclass

import legged_lab_v2.tasks.manager_based.locomotion.velocity.mdp as mdp
from legged_lab_v2.tasks.manager_based.locomotion.velocity.velocity_env_cfg import LocomotionVelocityRoughEnvCfg

##
# Pre-defined configs
##
from legged_lab_v2.assets.LW_leg import LW_LEG_CFG  # isort: skip
from legged_lab_v2.tasks.manager_based.locomotion.velocity.mdp.terrains.terrains_cfg import BLIND_ROUGH_TERRAINS_CFG, BLIND_ROUGH_AND_STAIRS_TERRAINS_CFG

@configclass
class LWLegBlindRoughEnvCfg(LocomotionVelocityRoughEnvCfg):
    base_link_name = "base_link"
    foot_link_name = ".*_foot_link"
    # fmt: off
    joint_names = [
        "right_hip_joint",
        "right_thigh_joint",
        "right_shank_joint",
        "right_wheel_joint",
        "right_foot_joint",
        "left_hip_joint",
        "left_thigh_joint",
        "left_shank_joint",
        "left_wheel_joint",
        "left_foot_joint"
    ]
    # fmt: on

    def __post_init__(self):
        # post init of parent
        super().__post_init__()

        # ------------------------------Sence------------------------------
        self.scene.robot = LW_LEG_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
        self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
        self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
        self.scene.terrain.terrain_generator = BLIND_ROUGH_AND_STAIRS_TERRAINS_CFG
        self.scene.terrain.max_init_terrain_level = 0

        # ------------------------------Observations------------------------------
        self.observations.policy.base_ang_vel.scale = 0.25
        self.observations.policy.joint_pos.scale = 1.0
        self.observations.policy.joint_vel.scale = 0.05
        self.observations.policy.height_scan = None
        self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joint_names
        self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joint_names

        # ------------------------------Actions------------------------------
        # reduce action scale
        self.actions.joint_pos.scale = 0.25
        self.actions.joint_pos.joint_names = self.joint_names
        self.actions.joint_pos.clip = {".*": (-100.0, 100.0)}

        # ------------------------------Events------------------------------
        self.events.randomize_rigid_body_mass_base.params["asset_cfg"].body_names = [self.base_link_name]
        self.events.randomize_rigid_body_mass_others.params["asset_cfg"].body_names = [
            "right_.*", "left_.*"
        ]
        self.events.randomize_com_positions.params["asset_cfg"].body_names = [self.base_link_name]
        self.events.randomize_reset_joints.params["velocity_range"] = (0.0, 0.0)
        self.events.randomize_apply_external_force_torque = None

        # ------------------------------Rewards------------------------------
        # General
        self.rewards.is_terminated.weight = -200.0
        # self.rewards.keep_alive = 1.0
        self.rewards.upward.weight = 1.0

        # Root penalties
        self.rewards.lin_vel_z_l2.weight = -0.5
        self.rewards.ang_vel_xy_l2.weight = -0.05 # -0.05
        self.rewards.flat_orientation_l2.weight = -5.0
        self.rewards.base_height_l2.weight = -100.0
        self.rewards.base_height_l2.params["target_height"] = 0.697 # 0.647
        self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name]

        # Joint penalties
        self.rewards.joint_torques_l2.weight = -1.5e-7 # 1.25e-5
        self.rewards.joint_acc_l2.weight = -1.25e-7
        self.rewards.create_joint_deviation_l1_rewterm("joint_deviation_up_l1", -0.15, [".*hip_joint", ".*wheel_joint"])
        self.rewards.create_joint_deviation_l1_rewterm(
            "joint_deviation_thigh_l1", -0.15, [".*thigh_joint"]
        )
        self.rewards.create_joint_deviation_l1_rewterm(
            "joint_deviation_ankle_l1", -0.2, [".*foot_joint"]
        )
        self.rewards.create_joint_deviation_l1_rewterm("joint_deviation_shank_l1", -0.15, [".*shank_joint"])
        # self.rewards.rew_keep_ankle_pitch_zero_in_air.weight = 0.5
        self.rewards.joint_pos_limits.weight = -2.0
        # self.rewards.joint_powers_l1.weight = -2e-5
        # self.rewards.stand_still.weight = -5.0
        self.rewards.joint_pos_penalty.weight = -2.0

        # Action penalties
        self.rewards.action_rate_l2.weight = -0.01 # -0.01
        # self.rewards.action_smoothness.weight = 0.0

        # Contact sensorstand_still
        self.rewards.undesired_contacts.weight = -10.0
        # self.rewards.contact_forces.weight = 0
        # self.rewards.contact_forces.params["sensor_cfg"].body_names = [self.foot_link_name]

        # Velocity-tracking rewards
        self.rewards.track_lin_vel_xy_exp.weight = 15.0
        self.rewards.track_lin_vel_xy_exp.func = mdp.track_lin_vel_xy_yaw_frame_exp
        self.rewards.track_ang_vel_z_exp.weight = 7.5
        self.rewards.track_ang_vel_z_exp.func = mdp.track_ang_vel_z_world_exp

        # Others
        self.rewards.bipedal_gait_reward.weight = 4.0
        self.rewards.feet_air_time.weight = 20.0
        self.rewards.feet_air_time.func = mdp.feet_air_time_positive_biped
        self.rewards.feet_air_time.params["threshold"] = 0.2
        self.rewards.feet_air_time.params["sensor_cfg"].body_names = [self.foot_link_name]
        # self.rewards.feet_contact.weight = 0
        # self.rewards.feet_contact.params["sensor_cfg"].body_names = [self.foot_link_name]
        # self.rewards.feet_contact_without_cmd.weight = 0
        # self.rewards.feet_contact_without_cmd.params["sensor_cfg"].body_names = [self.foot_link_name]
        # self.rewards.feet_stumble.weight = 0
        # self.rewards.feet_stumble.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.weight = -0.2
        self.rewards.feet_slide.params["sensor_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_slide.params["asset_cfg"].body_names = [self.foot_link_name]
        # self.rewards.feet_height.weight = 0
        # self.rewards.feet_height.params["target_height"] = 0.05
        # self.rewards.feet_height.params["asset_cfg"].body_names = [self.foot_link_name]
        # self.rewards.feet_height_body.weight = -20.0
        # self.rewards.feet_height_body.params["target_height"] = -0.346
        # self.rewards.feet_height_body.params["asset_cfg"].body_names = [self.foot_link_name]
        self.rewards.feet_distance_y_exp.weight = 5.0
        self.rewards.feet_distance_y_exp.params["stance_width"] = 0.42
        self.rewards.feet_distance_penalize.weight = -100.0
        self.rewards.feet_distance_penalize.params["min_feet_distance"] = 0.15

        # If the weight of rewards is 0, set rewards to None
        if self.__class__.__name__ == "LWLegBlindRoughEnvCfg":
            self.disable_zero_weight_rewards()

        # ------------------------------Terminations------------------------------
        self.terminations.illegal_contact.params["sensor_cfg"].body_names = [self.base_link_name]

        # ------------------------------Curriculums------------------------------
        # self.curriculum.command_levels.params["range_multiplier"] = (0.2, 1.0)
        self.curriculum.command_levels = None

        # ------------------------------Commands------------------------------
        self.commands.base_velocity.ranges.lin_vel_x = (-1.0, 1.0)
        self.commands.base_velocity.ranges.lin_vel_y = (-0.7, 0.7)
        self.commands.base_velocity.ranges.ang_vel_z = (-1.0, 1.0)
